package com.caseprocessing.enhanced;

import com.caseprocessing.ai.AIServiceManager;
import com.caseprocessing.api.DocumentAnalysisResult;
import com.caseprocessing.model.CaseDocument;
import com.caseprocessing.model.DocumentSegment;
import com.caseprocessing.model.DocumentType;
import com.caseprocessing.service.DocumentSplitter;

import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;

/**
 * AI增强的文档分割器
 * 结合传统算法和AI能力，实现智能文档分割
 */
public class AIEnhancedDocumentSplitter {
    
    private final DocumentSplitter traditionalSplitter;
    private final AIServiceManager aiServiceManager;
    private final SplitConfig config;
    
    public AIEnhancedDocumentSplitter(
            DocumentSplitter traditionalSplitter, 
            AIServiceManager aiServiceManager) {
        this.traditionalSplitter = traditionalSplitter;
        this.aiServiceManager = aiServiceManager;
        this.config = new SplitConfig();
    }
    
    /**
     * AI增强的文档分割主方法
     * 多模态分析 + 传统算法 + 智能融合
     */
    public CompletableFuture<List<DocumentSegment>> splitDocument(
            CaseDocument caseDocument, 
            List<String> documentImages) {
        
        // 1. 并行执行传统分割和AI分析
        CompletableFuture<List<DocumentSegment>> traditionalSplit = 
            CompletableFuture.supplyAsync(() -> traditionalSplitter.splitDocument(caseDocument));
        
        CompletableFuture<List<DocumentSegment>> aiEnhancedSplit = 
            aiServiceManager.aiEnhancedDocumentSplit(caseDocument, documentImages);
        
        // 2. 融合分割结果
        return CompletableFuture.allOf(traditionalSplit, aiEnhancedSplit)
                .thenCompose(v -> {
                    try {
                        List<DocumentSegment> traditional = traditionalSplit.get();
                        List<DocumentSegment> aiEnhanced = aiEnhancedSplit.get();
                        
                        return fuseSplitResults(caseDocument, traditional, aiEnhanced, documentImages);
                    } catch (Exception e) {
                        // AI分割失败时降级到传统分割
                        System.out.println("AI分割失败，降级到传统分割: " + e.getMessage());
                        return traditionalSplit;
                    }
                });
    }
    
    /**
     * 融合分割结果
     * 综合传统算法和AI分析的优势
     */
    private CompletableFuture<List<DocumentSegment>> fuseSplitResults(
            CaseDocument caseDocument,
            List<DocumentSegment> traditionalSegments,
            List<DocumentSegment> aiEnhancedSegments,
            List<String> documentImages) {
        
        // 1. 评估两种分割结果的质量
        SplitQuality traditionalQuality = evaluateSplitQuality(traditionalSegments);
        SplitQuality aiQuality = evaluateSplitQuality(aiEnhancedSegments);
        
        // 2. 选择最佳分割策略
        List<DocumentSegment> bestSegments;
        
        if (aiQuality.getOverallScore() > traditionalQuality.getOverallScore() + config.getAiThreshold()) {
            // AI分割明显更好
            bestSegments = aiEnhancedSegments;
            System.out.println("选择AI增强分割结果，质量评分: " + aiQuality.getOverallScore());
        } else if (traditionalQuality.getOverallScore() > aiQuality.getOverallScore() + config.getTraditionalThreshold()) {
            // 传统分割明显更好
            bestSegments = traditionalSegments;
            System.out.println("选择传统分割结果，质量评分: " + traditionalQuality.getOverallScore());
        } else {
            // 质量相近，进行智能融合
            bestSegments = intelligentFusion(traditionalSegments, aiEnhancedSegments);
            System.out.println("执行智能融合分割");
        }
        
        // 3. 对最终结果进行优化和验证
        return optimizeAndValidateSegments(bestSegments, caseDocument, documentImages);
    }
    
    /**
     * 智能融合两种分割结果
     */
    private List<DocumentSegment> intelligentFusion(
            List<DocumentSegment> traditionalSegments,
            List<DocumentSegment> aiEnhancedSegments) {
        
        List<DocumentSegment> fusedSegments = new ArrayList<>();
        
        // 使用区间合并算法融合分割点
        Set<Integer> allSplitPoints = new TreeSet<>();
        
        // 收集所有分割点
        for (DocumentSegment segment : traditionalSegments) {
            allSplitPoints.add(segment.getStartPage());
            allSplitPoints.add(segment.getEndPage() + 1); // 结束页的下一页作为分割点
        }
        
        for (DocumentSegment segment : aiEnhancedSegments) {
            allSplitPoints.add(segment.getStartPage());
            allSplitPoints.add(segment.getEndPage() + 1);
        }
        
        // 根据融合的分割点创建新的分段
        List<Integer> sortedSplitPoints = new ArrayList<>(allSplitPoints);
        
        for (int i = 0; i < sortedSplitPoints.size() - 1; i++) {
            int startPage = sortedSplitPoints.get(i);
            int endPage = sortedSplitPoints.get(i + 1) - 1;
            
            if (endPage >= startPage) {
                DocumentSegment fusedSegment = new DocumentSegment();
                fusedSegment.setStartPage(startPage);
                fusedSegment.setEndPage(endPage);
                fusedSegment.setConfidence(calculateFusedConfidence(startPage, endPage, traditionalSegments, aiEnhancedSegments));
                fusedSegment.setSplitMethod("智能融合");
                
                fusedSegments.add(fusedSegment);
            }
        }
        
        return fusedSegments;
    }
    
    /**
     * 计算融合片段的置信度
     */
    private double calculateFusedConfidence(
            int startPage, 
            int endPage,
            List<DocumentSegment> traditionalSegments,
            List<DocumentSegment> aiEnhancedSegments) {
        
        // 检查这个区间在两种分割中的支持程度
        double traditionalSupport = calculateSegmentSupport(startPage, endPage, traditionalSegments);
        double aiSupport = calculateSegmentSupport(startPage, endPage, aiEnhancedSegments);
        
        // 加权平均，AI支持度权重稍高
        return traditionalSupport * 0.4 + aiSupport * 0.6;
    }
    
    /**
     * 计算片段在分割结果中的支持度
     */
    private double calculateSegmentSupport(
            int startPage, 
            int endPage, 
            List<DocumentSegment> segments) {
        
        for (DocumentSegment segment : segments) {
            if (segment.getStartPage() == startPage && segment.getEndPage() == endPage) {
                return segment.getConfidence();
            }
            
            // 计算重叠度
            int overlapStart = Math.max(startPage, segment.getStartPage());
            int overlapEnd = Math.min(endPage, segment.getEndPage());
            
            if (overlapEnd >= overlapStart) {
                int overlapPages = overlapEnd - overlapStart + 1;
                int totalPages = endPage - startPage + 1;
                double overlapRatio = (double) overlapPages / totalPages;
                
                return segment.getConfidence() * overlapRatio;
            }
        }
        
        return 0.0;
    }
    
    /**
     * 优化和验证分割结果
     */
    private CompletableFuture<List<DocumentSegment>> optimizeAndValidateSegments(
            List<DocumentSegment> segments,
            CaseDocument caseDocument,
            List<String> documentImages) {
        
        // 1. 填充片段内容
        for (DocumentSegment segment : segments) {
            String content = extractSegmentContent(caseDocument, segment.getStartPage(), segment.getEndPage());
            segment.setContent(content);
            segment.setDocumentId(caseDocument.getDocumentId());
        }
        
        // 2. AI增强的文档类型识别
        List<CompletableFuture<Void>> classificationFutures = new ArrayList<>();
        
        for (int i = 0; i < segments.size(); i++) {
            DocumentSegment segment = segments.get(i);
            List<String> segmentImages = extractSegmentImages(documentImages, segment.getStartPage(), segment.getEndPage());
            
            CompletableFuture<Void> classificationFuture = 
                aiServiceManager.aiEnhancedDocumentClassification(segment, segmentImages)
                    .thenAccept(documentType -> {
                        segment.setDocumentType(documentType);
                        segment.setClassificationMethod("AI增强");
                    })
                    .exceptionally(throwable -> {
                        // 分类失败时使用传统方法
                        segment.setDocumentType(DocumentType.OTHER);
                        segment.setClassificationMethod("降级处理");
                        return null;
                    });
            
            classificationFutures.add(classificationFuture);
        }
        
        // 3. 等待所有分类完成
        return CompletableFuture.allOf(classificationFutures.toArray(new CompletableFuture[0]))
                .thenApply(v -> {
                    // 4. 最终验证和调整
                    return validateAndAdjustSegments(segments);
                });
    }
    
    /**
     * 验证和调整分割结果
     */
    private List<DocumentSegment> validateAndAdjustSegments(List<DocumentSegment> segments) {
        List<DocumentSegment> validatedSegments = new ArrayList<>();
        
        for (DocumentSegment segment : segments) {
            // 验证片段的合理性
            if (isValidSegment(segment)) {
                validatedSegments.add(segment);
            } else {
                // 处理无效片段
                System.out.println("发现无效片段: " + segment.getStartPage() + "-" + segment.getEndPage());
                
                // 尝试修复或分解片段
                List<DocumentSegment> repairedSegments = repairSegment(segment);
                validatedSegments.addAll(repairedSegments);
            }
        }
        
        // 确保片段连续性和完整性
        return ensureContinuity(validatedSegments);
    }
    
    /**
     * 评估分割质量
     */
    private SplitQuality evaluateSplitQuality(List<DocumentSegment> segments) {
        if (segments == null || segments.isEmpty()) {
            return new SplitQuality(0.0, 0.0, 0.0, 0.0);
        }
        
        // 1. 置信度评分
        double confidenceScore = segments.stream()
                .mapToDouble(DocumentSegment::getConfidence)
                .average()
                .orElse(0.0);
        
        // 2. 完整性评分
        double completenessScore = calculateCompleteness(segments);
        
        // 3. 一致性评分
        double consistencyScore = calculateConsistency(segments);
        
        // 4. 合理性评分
        double reasonablenessScore = calculateReasonableness(segments);
        
        return new SplitQuality(confidenceScore, completenessScore, consistencyScore, reasonablenessScore);
    }
    
    // ============ 辅助方法 ============
    
    private String extractSegmentContent(CaseDocument caseDocument, int startPage, int endPage) {
        // 实际实现中会从文档中提取指定页面的内容
        return "页面 " + startPage + " 到 " + endPage + " 的内容";
    }
    
    private List<String> extractSegmentImages(List<String> allImages, int startPage, int endPage) {
        if (allImages == null || allImages.isEmpty()) {
            return Collections.emptyList();
        }
        
        // 根据页面范围提取对应的图像
        List<String> segmentImages = new ArrayList<>();
        for (int page = startPage; page <= endPage && page - 1 < allImages.size(); page++) {
            segmentImages.add(allImages.get(page - 1));
        }
        
        return segmentImages;
    }
    
    private boolean isValidSegment(DocumentSegment segment) {
        return segment.getStartPage() > 0 && 
               segment.getEndPage() >= segment.getStartPage() &&
               segment.getContent() != null && 
               !segment.getContent().trim().isEmpty();
    }
    
    private List<DocumentSegment> repairSegment(DocumentSegment segment) {
        // 简单修复：如果片段过大，尝试分解
        List<DocumentSegment> repairedSegments = new ArrayList<>();
        
        int pageSpan = segment.getEndPage() - segment.getStartPage() + 1;
        if (pageSpan > config.getMaxPagesPerSegment()) {
            // 分解为较小的片段
            int segmentSize = config.getMaxPagesPerSegment();
            for (int start = segment.getStartPage(); start <= segment.getEndPage(); start += segmentSize) {
                int end = Math.min(start + segmentSize - 1, segment.getEndPage());
                
                DocumentSegment repairedSegment = new DocumentSegment();
                repairedSegment.setStartPage(start);
                repairedSegment.setEndPage(end);
                repairedSegment.setConfidence(segment.getConfidence() * 0.8); // 降低置信度
                repairedSegment.setSplitMethod("修复分解");
                
                repairedSegments.add(repairedSegment);
            }
        } else {
            repairedSegments.add(segment);
        }
        
        return repairedSegments;
    }
    
    private List<DocumentSegment> ensureContinuity(List<DocumentSegment> segments) {
        // 排序并检查连续性
        segments.sort(Comparator.comparingInt(DocumentSegment::getStartPage));
        
        List<DocumentSegment> continuousSegments = new ArrayList<>();
        
        for (int i = 0; i < segments.size(); i++) {
            DocumentSegment current = segments.get(i);
            
            if (i > 0) {
                DocumentSegment previous = continuousSegments.get(continuousSegments.size() - 1);
                
                // 检查是否有间隙
                if (current.getStartPage() > previous.getEndPage() + 1) {
                    // 填补间隙
                    DocumentSegment gap = new DocumentSegment();
                    gap.setStartPage(previous.getEndPage() + 1);
                    gap.setEndPage(current.getStartPage() - 1);
                    gap.setConfidence(0.3); // 低置信度
                    gap.setSplitMethod("间隙填补");
                    gap.setDocumentType(DocumentType.OTHER);
                    
                    continuousSegments.add(gap);
                }
            }
            
            continuousSegments.add(current);
        }
        
        return continuousSegments;
    }
    
    private double calculateCompleteness(List<DocumentSegment> segments) {
        // 检查页面覆盖的完整性
        Set<Integer> coveredPages = new HashSet<>();
        for (DocumentSegment segment : segments) {
            for (int page = segment.getStartPage(); page <= segment.getEndPage(); page++) {
                coveredPages.add(page);
            }
        }
        
        // 假设文档有100页
        int totalPages = 100;
        return (double) coveredPages.size() / totalPages;
    }
    
    private double calculateConsistency(List<DocumentSegment> segments) {
        // 检查分割的一致性
        double totalVariance = 0.0;
        for (int i = 1; i < segments.size(); i++) {
            int currentSize = segments.get(i).getEndPage() - segments.get(i).getStartPage() + 1;
            int previousSize = segments.get(i - 1).getEndPage() - segments.get(i - 1).getStartPage() + 1;
            totalVariance += Math.abs(currentSize - previousSize);
        }
        
        return Math.max(0.0, 1.0 - totalVariance / (segments.size() * 10.0));
    }
    
    private double calculateReasonableness(List<DocumentSegment> segments) {
        // 检查分割的合理性（片段大小、类型分布等）
        long reasonableSegments = segments.stream()
                .filter(segment -> {
                    int size = segment.getEndPage() - segment.getStartPage() + 1;
                    return size >= config.getMinPagesPerSegment() && size <= config.getMaxPagesPerSegment();
                })
                .count();
        
        return (double) reasonableSegments / segments.size();
    }
    
    /**
     * 分割质量评估结果
     */
    private static class SplitQuality {
        private final double confidenceScore;
        private final double completenessScore;
        private final double consistencyScore;
        private final double reasonablenessScore;
        
        public SplitQuality(double confidenceScore, double completenessScore, 
                           double consistencyScore, double reasonablenessScore) {
            this.confidenceScore = confidenceScore;
            this.completenessScore = completenessScore;
            this.consistencyScore = consistencyScore;
            this.reasonablenessScore = reasonablenessScore;
        }
        
        public double getOverallScore() {
            return (confidenceScore * 0.3 + completenessScore * 0.3 + 
                   consistencyScore * 0.2 + reasonablenessScore * 0.2);
        }
        
        public double getConfidenceScore() { return confidenceScore; }
        public double getCompletenessScore() { return completenessScore; }
        public double getConsistencyScore() { return consistencyScore; }
        public double getReasonablenessScore() { return reasonablenessScore; }
    }
    
    /**
     * 分割配置
     */
    private static class SplitConfig {
        private double aiThreshold = 0.1;
        private double traditionalThreshold = 0.1;
        private int minPagesPerSegment = 1;
        private int maxPagesPerSegment = 20;
        
        public double getAiThreshold() { return aiThreshold; }
        public double getTraditionalThreshold() { return traditionalThreshold; }
        public int getMinPagesPerSegment() { return minPagesPerSegment; }
        public int getMaxPagesPerSegment() { return maxPagesPerSegment; }
    }
}